{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp long_horizon2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Long-Horizon Original Datasets\n",
    "\n",
    "> Download and wrangling utility for long-horizon datasets. These datasets have been used by `NHITS, AutoFormer, Informer, PatchTST, TiDE` among many other neural forecasting methods. The datasets include the original [ETTh1, ETTh2, ETTm1, ETTm2, Weather, ILI, TrafficL](https://github.com/zhouhaoyi/ETDataset) benchmark datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| eval: false\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import os\n",
    "from dataclasses import dataclass\n",
    "from typing import Optional, Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "from datasetsforecast.utils import download_file, Info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@dataclass\n",
    "class ETTh1:\n",
    "    \"\"\"\n",
    "    The ETTh1 dataset monitors an electricity transformer \n",
    "    from a region of a province of China including oil temperature \n",
    "    and variants of load (such as high useful load and high useless load) \n",
    "    from July 2016 to July 2018 at an hourly frequency.\n",
    "    \"\"\"\n",
    "    freq: str = 'H'\n",
    "    name: str = 'ETTh1'\n",
    "    n_ts: int = 7\n",
    "    n_time: int = 14_400\n",
    "    test_size: int = 2_880\n",
    "    val_size: int = 2_880\n",
    "    horizons: Tuple[int] = (96, 192, 336, 720)\n",
    "        \n",
    "@dataclass\n",
    "class ETTh2:\n",
    "    \"\"\"\n",
    "    The ETTh2 dataset monitors an electricity transformer \n",
    "    from a region of a province of China including oil temperature \n",
    "    and variants of load (such as high useful load and high useless load) \n",
    "    from July 2016 to July 2018 at an hourly frequency.\n",
    "    \"\"\"    \n",
    "    freq: str = 'H'\n",
    "    name: str = 'ETTh2'\n",
    "    n_ts: int = 7\n",
    "    n_time: int = 14_400\n",
    "    test_size: int = 2_880\n",
    "    val_size: int = 2_880\n",
    "    horizons: Tuple[int] = (96, 192, 336, 720)\n",
    "\n",
    "@dataclass\n",
    "class ETTm1:\n",
    "    \"\"\"\n",
    "    The ETTm1 dataset monitors an electricity transformer \n",
    "    from a region of a province of China including oil temperature \n",
    "    and variants of load (such as high useful load and high useless load) \n",
    "    from July 2016 to July 2018 at a fifteen minute frequency.\n",
    "    \"\"\"    \n",
    "    freq: str = '15T'\n",
    "    name: str = 'ETTm1'\n",
    "    n_ts: int = 7\n",
    "    n_time: int = 57_600\n",
    "    test_size: int = 11_520\n",
    "    val_size: int = 11_520\n",
    "    horizons: Tuple[int] = (96, 192, 336, 720)\n",
    "        \n",
    "@dataclass\n",
    "class ETTm2:\n",
    "    \"\"\"\n",
    "    The ETTm2 dataset monitors an electricity transformer \n",
    "    from a region of a province of China including oil temperature \n",
    "    and variants of load (such as high useful load and high useless load) \n",
    "    from July 2016 to July 2018 at a fifteen minute frequency.\n",
    "    \n",
    "        Reference:\n",
    "        Zhou, et al. Informer: Beyond Efficient Transformer \n",
    "        for Long Sequence Time-Series Forecasting. AAAI 2021.\n",
    "        https://arxiv.org/abs/2012.07436\n",
    "    \"\"\"\n",
    "    freq: str = '15T'\n",
    "    name: str = 'ETTm2'\n",
    "    n_ts: int = 7\n",
    "    n_time: int = 57_600\n",
    "    test_size: int = 11_520\n",
    "    val_size: int = 11_520\n",
    "    horizons: Tuple[int] = (96, 192, 336, 720)\n",
    "\n",
    "@dataclass\n",
    "class ECL:\n",
    "    \"\"\"\n",
    "    The Electricity dataset reports the fifteen minute electricity \n",
    "    consumption (KWh) of 321 customers from 2012 to 2014. \n",
    "    For comparability, we aggregate it hourly.\n",
    "    \n",
    "        Reference:\n",
    "        Li, S et al. Enhancing the locality and breaking the memory \n",
    "        bottleneck of Transformer on time series forecasting.\n",
    "        NeurIPS 2019. http://arxiv.org/abs/1907.00235.\n",
    "    \"\"\"\n",
    "    freq: str = '15T'\n",
    "    name: str = 'ECL'\n",
    "    n_ts: int = 321\n",
    "    n_time: int = 26_304\n",
    "    test_size: int = 5_260\n",
    "    val_size: int = 2_632\n",
    "    horizons: Tuple[int] = (96, 192, 336, 720)\n",
    "\n",
    "@dataclass\n",
    "class TrafficL:\n",
    "    \"\"\"\n",
    "    This large Traffic dataset was collected by the California Department \n",
    "    of Transportation, it reports road hourly occupancy rates of 862 sensors, \n",
    "    from January 2015 to December 2016.\n",
    "    \n",
    "        Reference:\n",
    "        Lai, G., Chang, W., Yang, Y., and Liu, H. Modeling Long and\n",
    "        Short-Term Temporal Patterns with Deep Neural Networks.\n",
    "        SIGIR 2018. http://arxiv.org/abs/1703.07015.\n",
    "        \n",
    "        Wu, H., Xu, J., Wang, J., and Long, M. Autoformer:\n",
    "        Decomposition Transformers with auto-correlation for\n",
    "        long-term series forecasting. NeurIPS 2021. \n",
    "        https://arxiv.org/abs/2106.13008.        \n",
    "    \"\"\"\n",
    "    freq: str = 'H'\n",
    "    name: str = 'traffic'\n",
    "    n_ts: int = 862\n",
    "    n_time = 17_544\n",
    "    test_size: int = 3_508\n",
    "    val_size: int = 1_756\n",
    "    horizons: Tuple[int] = (96, 192, 336, 720)\n",
    "\n",
    "@dataclass\n",
    "class Weather:\n",
    "    \"\"\"\n",
    "    This Weather dataset contains the 2020 year of 21 meteorological \n",
    "    measurements\n",
    "    recorded every 10 minutes from the Weather Station of the Max Planck Biogeochemistry \n",
    "    Institute in Jena, Germany.\n",
    "\n",
    "        Reference:\n",
    "        Wu, H., Xu, J., Wang, J., and Long, M. Autoformer:\n",
    "        Decomposition Transformers with auto-correlation for\n",
    "        long-term series forecasting. NeurIPS 2021. \n",
    "        https://arxiv.org/abs/2106.13008.\n",
    "    \"\"\"\n",
    "    freq: str = '10M'\n",
    "    name: str = 'weather'\n",
    "    n_ts: int = 21\n",
    "    n_time = 52_695\n",
    "    test_size: int = 10_539\n",
    "    val_size: int = 5_270\n",
    "    horizons: Tuple[int] = (96, 192, 336, 720)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "LongHorizon2Info = Info((\n",
    "        ETTh1, ETTh2, ETTm1, ETTm2, \n",
    "        ECL, TrafficL, Weather\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@dataclass\n",
    "class LongHorizon2:\n",
    "    \"\"\"\n",
    "    This Long-Horizon datasets wrapper class, provides\n",
    "    with utility to download and wrangle the following datasets:    \n",
    "    ETT, ECL, Exchange, Traffic, ILI and Weather.\n",
    "    \n",
    "    - Each set is normalized with the train data mean and standard deviation.\n",
    "    - Datasets are partitioned into train, validation and test splits.\n",
    "    - For all datasets: 70%, 10%, and 20% of observations are train, validation, test, \n",
    "      except ETT that uses 20% validation.  \n",
    "    \"\"\"\n",
    "\n",
    "    source_url: str = 'https://www.dropbox.com/s/rlc1qmprpvuqrsv/all_six_datasets.zip?dl=1'\n",
    "    #source_url: str = 'https://drive.google.com/file/d/1alE33S1GmP5wACMXaLu50rDIoVzBM4ik/view?usp=share_link'\n",
    "    # See https://github.com/google-research/google-research/blob/master/tide/\n",
    "    \n",
    "    @staticmethod\n",
    "    def _normalize_data(data_mat, test_size):\n",
    "        \"\"\"\n",
    "        Normalizes dataset.\n",
    "\n",
    "            Parameters\n",
    "            ----------     \n",
    "            data_mat: np.array\n",
    "                panel data in matrix of shape (n_time, n_series)\n",
    "            test_size: int\n",
    "                observations kept as test.\n",
    "        \"\"\"\n",
    "        scaler = StandardScaler()\n",
    "        train_mat = data_mat[:-test_size, :]\n",
    "        scaler = scaler.fit(train_mat)\n",
    "        data_mat = scaler.transform(data_mat)\n",
    "        return data_mat\n",
    "\n",
    "    @staticmethod\n",
    "    def load(directory: str,\n",
    "             group: str,\n",
    "             normalize: bool=True) -> pd.DataFrame:\n",
    "        \"\"\"\n",
    "\n",
    "        Downloads and long-horizon forecasting benchmark datasets.\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            directory: str\n",
    "                Directory where data will be downloaded.\n",
    "            group: str\n",
    "                Group name.\n",
    "                Allowed groups: 'ETTh1', 'ETTh2', \n",
    "                                'ETTm1', 'ETTm2',\n",
    "                                'ECL', 'Exchange',\n",
    "                                'Traffic', 'Weather', 'ILI'.\n",
    "            normalize: bool\n",
    "                If `True` std. normalize data or not\n",
    "\n",
    "            Returns\n",
    "            ------- \n",
    "            y_df: pd.DataFrame\n",
    "                Target time series with columns ['unique_id', 'ds', 'y'].\n",
    "        \"\"\"\n",
    "        if group not in LongHorizon2Info.groups:\n",
    "            raise Exception(f'group not found {group}')\n",
    "\n",
    "        LongHorizon2.download(directory)\n",
    "        path = f'{directory}/longhorizon2/all_six_datasets'\n",
    "\n",
    "        # Read and parse data\n",
    "        Y_df = pd.read_csv(f'{path}/{group}/Y_df.csv')\n",
    "\n",
    "        Y_df = Y_df.set_index('date')\n",
    "        Y_df = Y_df.melt(ignore_index=False).reset_index()\n",
    "        Y_df = Y_df.rename(columns={'date': 'ds', 'variable': 'unique_id',\n",
    "                                    'value': 'y'})\n",
    "        Y_df['ds'] = pd.to_datetime(Y_df['ds'])\n",
    "\n",
    "        # Mimic weird data drop following Google\n",
    "        # https://github.com/google-research/google-research/blob/master/tide/train.py#L78\n",
    "        # https://github.com/google-research/google-research/blob/master/tide/data_loader.py#L184\n",
    "        n_time = LongHorizon2Info[group].n_time\n",
    "        Y_df = Y_df.groupby('unique_id').head(n_time).reset_index()\n",
    "\n",
    "        if normalize:\n",
    "            n_ts = LongHorizon2Info[group].n_ts\n",
    "            test_size = LongHorizon2Info[group].test_size\n",
    "            data_mat  = Y_df.y.values.reshape(n_ts, -1).transpose()\n",
    "            \n",
    "            data_mat  = LongHorizon2._normalize_data(data_mat=data_mat, \n",
    "                                                     test_size=test_size)\n",
    "\n",
    "            data_mat = data_mat.transpose().flatten()\n",
    "            Y_df['y'] = data_mat\n",
    "\n",
    "        return Y_df\n",
    "\n",
    "    @staticmethod\n",
    "    def download(directory: str) -> None:\n",
    "        \"\"\"\n",
    "        Download Long Horizon 2 Datasets.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        directory: str\n",
    "            Directory path to download dataset.\n",
    "        \"\"\"\n",
    "        path = f'{directory}/longhorizon2/'\n",
    "        if not os.path.exists(path):\n",
    "             download_file(path, LongHorizon2.source_url, decompress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| eval: false\n",
    "# Checking qualitatively that series are correctly normalized\n",
    "group = 'ETTh1'\n",
    "Y_df1 = LongHorizon2.load(directory='./data', group=group)\n",
    "Y_df2 = LongHorizon2.load(directory='./data', group=group, normalize=False)\n",
    "\n",
    "unique_id = 'OT'\n",
    "\n",
    "plot_df1 = Y_df1[Y_df1.unique_id==unique_id]\n",
    "plot_df2 = Y_df2[Y_df2.unique_id==unique_id]\n",
    "\n",
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "ax2 = ax1.twinx()\n",
    "#ax1.plot(plot_df1.y[-2000:], 'b-')\n",
    "#ax1.plot((plot_df2.y[-2000:]-38)/10, 'g-')\n",
    "\n",
    "ax1.plot(plot_df1.y[-2000:], 'b-')\n",
    "ax2.plot(plot_df2.y[-2000:], 'g-')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unit testing for abscense of duplicate unique_id-ds\n",
    "# Unit testing for correct number of series\n",
    "# Unit testing for correct number of date stamps\n",
    "for group, meta in LongHorizon2Info:\n",
    "    Y_df = LongHorizon2.load(directory='data', group=group)\n",
    "    unique_elements = Y_df.groupby(['unique_id', 'ds']).size()\n",
    "    unique_ts = Y_df.groupby('unique_id').size()\n",
    "    n_time = len(Y_df.ds.unique())\n",
    "\n",
    "    assert (unique_elements != 1).sum() == 0, f'Duplicated records found: {group}'\n",
    "    assert unique_ts.shape[0] == meta.n_ts, f'Number of time series not match: {group}'\n",
    "    assert n_time == meta.n_time, f'Number of time observations not match: {group} {n_time} { meta.n_time}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
